#!/usr/bin/env python
# coding=utf-8

# --------------------------------------------------------
# R-FCN
# Copyright (c) 2016 Yuwen Xiong
# Licensed under The MIT License [see LICENSE for details]
# Written by Yuwen Xiong
# --------------------------------------------------------

"""
Demo script showing detections in sample images.

See README.md for installation instructions before running.
"""
import sys

import _init_paths
from fast_rcnn.config import cfg
from fast_rcnn.test import im_detect
from fast_rcnn.nms_wrapper import nms
from utils.timer import Timer
import matplotlib.pyplot as plt
import xml.etree.ElementTree as ET
import numpy as np
import scipy.io as sio
import caffe, os, sys, cv2
import argparse
import shutil
import MySQLdb

CLASSES = (
'__background__', 'box')

prototxt = os.path.join(cfg.MODELS_DIR, 'ResNet-50', 'rfcn_end2end', 'test_agnostic.prototxt')

caffemodel = os.path.join(cfg.DATA_DIR, 'rfcn_models', 'resnet50_rfcn_ohem_iter_10000_shelf.caffemodel')
testPath = os.path.join('/home/letitia/py-R-FCN/data/VOCdevkit0712/VOC0712',"<0.5")
xmlPath = os.path.join('/home/letitia/py-R-FCN/data/VOCdevkit0712/VOC0712', 'xmlfile')

def vis_detections(im, class_name, dets, thresh):
    """Draw detected bounding boxes."""
    inds = np.where(dets[:, -1] >= thresh)[0]
    if len(inds) == 0:
        return

    im = im[:, :, (2, 1, 0)]
    fig, ax = plt.subplots(figsize=(12, 12))
    ax.imshow(im, aspect='equal')
    for i in inds:
        bbox = dets[i, :4]
        score = dets[i, -1]

        ax.add_patch(
            plt.Rectangle((bbox[0], bbox[1]),
                          bbox[2] - bbox[0],
                          bbox[3] - bbox[1], fill=False,
                          edgecolor='red', linewidth=3.5)
        )
        ax.text(bbox[0], bbox[1] - 2,
                '{:s} {:.3f}'.format(class_name, score),
                bbox=dict(facecolor='blue', alpha=0.5),
                fontsize=14, color='white')

    ax.set_title(('{} detections with '
                  'p({} | box) >= {:.1f}').format(class_name, class_name,
                                                  thresh),
                 fontsize=14)
    plt.axis('off')
    plt.tight_layout()
    plt.draw()


def buildNewsXmlFile(im_name, outputlist):
    anno = ET.Element("annotation")
    folder = ET.SubElement(anno,"folder")
    filename = ET.SubElement(anno, "filename")
    source = ET.SubElement(anno, "source")
    size = ET.SubElement(anno, "size")
    w = ET.SubElement(size, "width")
    h = ET.SubElement(size, "height")

    for i in range(len(outputlist)):
        object = ET.SubElement(anno, "object")
        o_name = ET.SubElement(object, "name")
        o_name.text = "box"
        o_class = ET.SubElement(object, "class")
        o_class.text = "others"
        o_bndbox = ET.SubElement(object, "bndbox")
        xmin = ET.SubElement(o_bndbox, "xmin")
        xmin.text = str(outputlist[i][0])
        ymin = ET.SubElement(o_bndbox, "ymin")
        ymin.text = str(outputlist[i][1])
        xmax = ET.SubElement(o_bndbox, "xmax")
        xmax.text = str(outputlist[i][2])
        ymax = ET.SubElement(o_bndbox, "ymax")
        ymax.text = str(outputlist[i][3])

    # 将节点数信息保存在ElementTree中，并且保存为XML格式文件
    tree = ET.ElementTree(anno)
    currentPath = os.path.join(xmlPath,im_name.split('.')[0]+".xml")
    f = open(currentPath, 'w')
    tree.write(f, encoding="UTF-8", xml_declaration="xml version='1.0'", method="xml")
    f.close()

def demo(net, image_name):
    """Detect object classes in an image using pre-computed object proposals."""

    # Load the demo image
    im_file = os.path.join(testPath, image_name)
    im = cv2.imread(im_file)

    # Detect all object classes and regress object bounds
    timer = Timer()
    timer.tic()
    print im_file
    try:
        scores, boxes = im_detect(net, im)
    except:
        print 'error'
        return
    timer.toc()
    print ('Detection took {:.3f}s for '
           '{:d} object proposals').format(timer.total_time, boxes.shape[0])

    # Visualize detections for each class
    CONF_THRESH = 0.7
    NMS_THRESH = 0.3
    detections = {}
    for cls_ind, cls in enumerate(CLASSES[1:]):
        cls_ind += 1  # because we skipped background
        cls_boxes = boxes[:, 4:8]
        cls_scores = scores[:, cls_ind]
        dets = np.hstack((cls_boxes,
                          cls_scores[:, np.newaxis])).astype(np.float32)
        keep = nms(dets, NMS_THRESH)
        dets = dets[keep, :]
        detections[cls] = dets.tolist()
        vis_detections(im, cls, dets, thresh=CONF_THRESH)

        # if detections is not None:
        #     outputlist = []
        #     for cls in detections:
        #         for box in detections[cls]:
        #             if box[4] > CONF_THRESH:
        #                 for k in range(0,4):
        #                     box[k] = round(box[k])
        #                 outputlist.append(box)
        #
        #     for i in range(len(outputlist)):
        #         for j in range(len(outputlist)):
        #             if outputlist[i][0] < outputlist[j][0]:
        #                 temp = outputlist[j]
        #                 outputlist[j] = outputlist[i]
        #                 outputlist[i] = temp
        #
        #     buildNewsXmlFile(image_name, outputlist)



def parse_args():
    """Parse input arguments."""
    parser = argparse.ArgumentParser(description='Faster R-CNN demo')

    args = parser.parse_args()

    return args


if __name__ == '__main__':
    cfg.TEST.HAS_RPN = True  # Use RPN for proposals

    args = parse_args()

    if not os.path.isfile(caffemodel):
        raise IOError(('{:s} not found.\n').format(caffemodel))

    caffe.set_mode_gpu()
    caffe.set_device(0)
    cfg.GPU_ID = 0

    net = caffe.Net(prototxt, caffemodel, caffe.TEST)

    print '\n\nLoaded network {:s}'.format(caffemodel)

    # Warmup on a dummy image
    im = 128 * np.ones((300, 500, 3), dtype=np.uint8)
    for i in xrange(2):
        _, _ = im_detect(net, im)

    files = os.listdir(testPath)
    for im_name in files:
        print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
        demo(net, im_name)

    plt.show()
